package com.sunny.zkdislock;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

public class ZkDisLock {

    private Logger log = LoggerFactory.getLogger(ZkDisLock.class);

    private static final String DEFAULT_HOST = "192.168.163.128";
    private static final int DEFAULT_PORT = 2181;

    private CountDownLatch countDownLatch = new CountDownLatch(1);
    private CountDownLatch locked = new CountDownLatch(1);
    private String host;
    private int port;
    private ZooKeeper zk;

    private String node;

    public ZkDisLock(String host, int port) throws IOException {
        this.host = host;
        this.port = port;
        this.zk = new ZooKeeper(host+":"+port, 20000,
                watchedEvent -> {
                    log.info("初始化成功");
                    try {
                        String res = "已存在";
                        if(!this.exists("/lock"))
                            res = zk.create("/lock","".getBytes(),ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
                        log.info("创建lock节点结果：{}",res);
                    } catch (KeeperException e) {
                        e.printStackTrace();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                    countDownLatch.countDown();
                });
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public void lock() throws InterruptedException, KeeperException, UnknownHostException {
        if (!tryLock()){
            locked.await();
        }
    }

    public boolean tryLock() throws UnknownHostException, KeeperException, InterruptedException {
        log.info("尝试获取锁");
        String path = "/lock/s";

        //可重入处理
        List<String> locks = getLocks();
        for(String lock: locks){
            String tmpLock = "/lock/"+lock;
//            System.out.println(getName()+" " + getData(tmpLock));
            if(getName().equals(getData(tmpLock))) {
                return judge(tmpLock);
            }
        }

        //不存在节点则创建等待节点
        node = zk.create(path, getName().getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL_SEQUENTIAL);

        //判断是否已经获得锁并加入watcher
        return judge(node);
    }

    public void unlock() throws KeeperException, InterruptedException, UnknownHostException {
        List<String> locks = getLocks();
        System.out.println(node + " " + ("/lock/"+locks.get(0)));
        if(locks.size() <= 0 || !node.equals("/lock/"+locks.get(0))){
            return;
        }
        log.info("释放锁");
        zk.delete("/lock/"+locks.get(0), -1);
    }

    public boolean exists(String path) throws KeeperException, InterruptedException {
        Stat stat = zk.exists(path, null);
        if(null == stat)
            return false;
        return true;
    }

    public List<String> getLocks() throws KeeperException, InterruptedException {
        List<String> locks = zk.getChildren("/lock",
                null);
        return locks;
    }

    public boolean judge(String res) throws KeeperException, InterruptedException{
        List<String> locks = getLocks();
        Collections.sort(locks);
        if(("/lock/"+locks.get(0)).equals(res)){
            locked.countDown();
            return true;
        }
        //加入watcher
        zk.exists("/lock", watchedEvent -> {
            try {
                tryLock();
            } catch (UnknownHostException e) {
                e.printStackTrace();
            } catch (KeeperException e) {
                e.printStackTrace();
            } catch (InterruptedException e) {
                e.printStackTrace();
            }
        });
        return false;
    }

    public String getData(String path) throws KeeperException, InterruptedException {
        String res = null;
        if(exists(path))
            res = new String(zk.getData(path, null, null));
        return res;
    }

    public String getName() throws UnknownHostException {
        InetAddress addr = InetAddress.getLocalHost();
        String tname = Thread.currentThread().getName();
        return addr.getHostAddress()+" "+tname;
    }

    public static void main(String[] args) {

        ExecutorService tp = Executors.newFixedThreadPool(5);

        try {
            ZkDisLock zkDisLock = new ZkDisLock(DEFAULT_HOST, DEFAULT_PORT);

            for(int i=0;i<5;i++) {
                tp.execute(() -> {
                    try {
                        zkDisLock.lock();
                        //模拟处理
                        System.out.println(Thread.currentThread().getName() + "获得锁，任务处理中...");
                        Thread.sleep(5000);
                        zkDisLock.unlock();
                    } catch (UnknownHostException e) {
                        e.printStackTrace();
                    } catch (KeeperException e) {
                        e.printStackTrace();
                    } catch (InterruptedException e) {
                        e.printStackTrace();
                    }
                });
            }
        } catch (IOException e) {
            e.printStackTrace();
        }


    }

}
